;;  this is x64 Windows Assembly

PUBLIC _penter
EXTRN	ThunkSetCurrentFunctionAndGetFunctionJumpAddress:PROC
EXTRN	ThunkShouldStubFunction:PROC	
EXTRN	GetCallingFunctionContext:PROC
EXTRN   IsMockerAlive:PROC

;-----------------------------------------Helper Macros---------------------------------
SaveArguments MACRO start

	mov [rsp+start+68h],rsi
	.savereg rsi, start+68h
	mov [rsp+start+60h],rdi 
	.savereg rdi, start+60h

	; Save integer arguments
	mov [rsp+start+58h],r9
	.savereg r9, start+58h	
    mov [rsp+start+50h],r8 
    .savereg r8, start+50h    
    mov [rsp+start+48h],rdx
    .savereg rdx, start+48h   
    mov [rsp+start+40h],rcx  
    .savereg rcx, start+40h
    
    ; Save FP arguments
	movups [rsp+start+30h], xmm3
    .savexmm128 xmm3, start+30h
    movups [rsp+start+20h], xmm2
    .savexmm128 xmm2, start+20h
    movups [rsp+start+10h], xmm1
    .savexmm128 xmm1, start+10h
    movups [rsp+start], xmm0	
    .savexmm128 xmm0, start
ENDM

RestoreArguments MACRO start
	mov rsi, [rsp+start+68h]
	mov rdi, [rsp+start+60h]

	; Restore integer arguments
	mov r9,  [rsp+start+58h]
	mov r8,  [rsp+start+50h]
	mov rdx, [rsp+start+48h]
	mov rcx, [rsp+start+40h]
	
	; Restore FP arguments
	movups xmm3,  [rsp+start+30h]
	movups xmm2,  [rsp+start+20h]
	movups xmm1,  [rsp+start+10h]
	movups xmm0,  [rsp+start]
ENDM

Process MACRO
	
	; If Mocker is not enable then we don't need to do anything
	call    IsMockerAlive
	movzx	eax, al
	test	eax, eax
	je		epilogue	

    ; Get calling function address & value of RSP right after calling this function.

	; The return value of GetCallingFunctionContext is larger than 64 bits
	; so a secret first parameter is passed which contains the address where the return value should be stored.
    mov		r9, rsp
    add		r9, 110h
    mov		rcx, r9		
    mov		rdx, 2		; Set stack trace depth.	
	call    GetCallingFunctionContext

	; GetCallingFunctionContext returns the following structure.
	; struct cFunctionContext
	; {
	;		size_t address;
	;		size_t RSP;
	; }
	; Read values from the structure returned by GetCallingFunctionContext.
	mov		r9,  [rsp+110h]	; Move RSP in r9.
	mov		r10, [rsp+118h]	; Move address in r10.
	mov		[rsp+108h], r10 ; We will need r10 later so we save its value on the stack.
	
	; Setup parameters before calling ThunkShouldStubFunction
	mov		QWORD PTR callingFunctionAddress$[rsp], r9
	mov	    rdx, [rsp+70h]	; If we are dealing with a member function then [rsp+70h] contains the this pointer.
	mov	    rcx, QWORD PTR callingFunctionAddress$[rsp]
	call	ThunkShouldStubFunction
	
	; If ThunkShouldStubFunction returned false then we don't do anything.
	movzx	eax, al
	test	eax, eax
	je		epilogue	
	
	; Get function jump address.
	; ThunkSetCurrentFunctionAndGetFunctionJumpAddress will return the jump address in Rax.
	mov	 rcx, QWORD PTR callingFunctionAddress$[rsp]
	call ThunkSetCurrentFunctionAndGetFunctionJumpAddress
			
	; Restore registers before calling mock function
	RestoreArguments 30h
	
	; Now jump to the mock function.
	mov rsp, [rsp+108h]	; Restore the value of Rsp that we previously sotred on the stack.
	jmp rax				; As explained above, the jump address is stored in Rax.
	
ENDM
;---------------------------------------------------------------------------------------
       
_TEXT SEGMENT ;WORD 'CODE'
_penter PROC
; Line 21
$LN4:
	; Save arguments
    mov r11, rdx
    mov r10, r9
    
    ; Check whether _penter was called from a leaf function.
    ; If _penter was called from a leaf function then the stack will be 16-byte aligned.
    xor edx, edx   
    mov r9,  16         
    mov rax, rsp		; Rsp contains the stack pointer .   
    div r9				; Divide rax by r9.
    cmp rdx, 8			; The remainder of the previous division was put in rdx.
    jne leafFunction    ; If the remainder is 8 then we are dealing with a non-leaf function.

	; Restore arguments
    mov rdx, r11
    mov r9, r10
	jmp _penterNonLeafFunction
    
leafFunction:
	; Restore arguments
    mov rdx, r11
    mov r9, r10
	jmp _penterLeafFunction


	ret					; This instruction is deliberately unreachable.
_penter ENDP
_TEXT   ENDS


PUBLIC _penterNonLeafFunction
_TEXT SEGMENT ;WORD 'CODE'
targetPointer$ = 8
callingFunctionAddress$ = 16	
_penterNonLeafFunction PROC FRAME    
    sub	rsp, 138h
	.allocstack 138h
	
	; Save arguments
    SaveArguments 30h
	
.endprolog

	Process
	
epilogue:
	RestoreArguments 30h
	add	rsp, 138h
	ret	0
_penterNonLeafFunction ENDP
_TEXT   ENDS  
 

 PUBLIC _penterLeafFunction   
_TEXT SEGMENT ;WORD 'CODE'
targetPointer$ = 8
callingFunctionAddress$ = 16
_penterLeafFunction PROC FRAME   
    sub	rsp, 130h
	.allocstack 130h
	
	; Save arguments
    SaveArguments 30h
	
.endprolog
	
	Process

epilogue:
	RestoreArguments 30h
	add rsp, 130h
	ret	0
_penterLeafFunction ENDP
_TEXT   ENDS
END
        

